import pandas as pd
from io import BytesIO
import os
from openpyxl.styles import Border, Side, Alignment, Font
from openpyxl import load_workbook

class ExcelUtils:
    @staticmethod
    def excel_export(data, headers, multi_level=True, sheet_name='Sheet1'):
        """
        将数据导出为 Excel 流
        :param data: 要导出的数据，列表的列表
        :param headers: 表头，可以是单级表头或多级表头
        :param multi_level: 是否为多级表头
        :param sheet_name: 工作表名称
        :return: Excel 文件的二进制流
        """
        output = BytesIO()
        try:
            with pd.ExcelWriter(output, engine='openpyxl') as writer:
                # 确保工作表存在
                if sheet_name not in writer.sheets:
                    writer.book.create_sheet(sheet_name)

                # 定义表头样式：居中且加粗
                header_style = Alignment(horizontal='center', vertical='center')
                bold_font = Font(bold=True)

                if multi_level:
                    df = pd.DataFrame(columns=pd.MultiIndex.from_tuples(headers))
                    # 获取多级表头
                    header = df.columns
                    num_levels = len(header.levels)

                    # 写入多级表头并合并单元格
                    for level in range(num_levels):
                        row_values = [cell[level] for cell in header]
                        writer.sheets[sheet_name].append(row_values)

                        # 合并当前层级中相邻且内容相同的单元格
                        start_col = 1
                        current_value = row_values[0]
                        for col, value in enumerate(row_values[1:], start=2):
                            if value != current_value:
                                if col - start_col > 1:
                                    writer.sheets[sheet_name].merge_cells(
                                        start_row=writer.sheets[sheet_name].max_row,
                                        start_column=start_col,
                                        end_row=writer.sheets[sheet_name].max_row,
                                        end_column=col - 1
                                    )
                                start_col = col
                                current_value = value

                        if len(row_values) - start_col + 1 > 1:
                            writer.sheets[sheet_name].merge_cells(
                                start_row=writer.sheets[sheet_name].max_row,
                                start_column=start_col,
                                end_row=writer.sheets[sheet_name].max_row,
                                end_column=len(row_values)
                            )

                        # 设置表头样式
                        for cell in writer.sheets[sheet_name].iter_cols(
                                min_row=writer.sheets[sheet_name].max_row,
                                max_row=writer.sheets[sheet_name].max_row
                        ):
                            for c in cell:
                                c.alignment = header_style
                                c.font = bold_font

                    # 合并不同层级间的单元格
                    for col in range(1, len(headers) + 1):
                        start_row = 1
                        current_value = writer.sheets[sheet_name].cell(row=start_row, column=col).value
                        for row in range(2, num_levels + 1):
                            cell_value = writer.sheets[sheet_name].cell(row=row, column=col).value
                            if cell_value != current_value:
                                if row - start_row > 1:
                                    writer.sheets[sheet_name].merge_cells(
                                        start_row=start_row,
                                        start_column=col,
                                        end_row=row - 1,
                                        end_column=col
                                    )
                                start_row = row
                                current_value = cell_value
                        if num_levels - start_row + 1 > 1:
                            writer.sheets[sheet_name].merge_cells(
                                start_row=start_row,
                                start_column=col,
                                end_row=num_levels,
                                end_column=col
                            )

                    if data is not None and len(data) > 0:
                        # 将数组对象转换为列表的列表
                        new_data = []
                        for item in data:
                            row = []
                            for header in headers:
                                if isinstance(header, tuple):
                                    key = header[-1]
                                else:
                                    key = header
                                row.append(item.get(key, ''))
                            new_data.append(row)

                        header_count = len(headers)
                        final_data = []
                        for row in new_data:
                            # 若数据长度大于表头长度，截取与表头长度相同的数据
                            if len(row) > header_count:
                                final_data.append(row[:header_count])
                            # 若数据长度小于表头长度，用空字符串补齐
                            elif len(row) < header_count:
                                new_row = row + [''] * (header_count - len(row))
                                final_data.append(new_row)
                            else:
                                final_data.append(row)
                        df = pd.DataFrame(final_data, columns=pd.MultiIndex.from_tuples(headers))
                        # 写入数据
                        for row in df.values.tolist():
                            writer.sheets[sheet_name].append(row)
                else:
                    df = pd.DataFrame(columns=headers)
                    df.to_excel(writer, sheet_name=sheet_name, index=False)
                    # 设置单级表头样式
                    for cell in writer.sheets[sheet_name].iter_cols(min_row=1, max_row=1):
                        for c in cell:
                            c.alignment = header_style
                            c.font = bold_font

                    if data is not None and len(data) > 0:
                        # 将数组对象转换为列表的列表
                        new_data = []
                        for item in data:
                            row = []
                            for header in headers:
                                row.append(item.get(header, ''))
                            new_data.append(row)

                        header_count = len(headers)
                        final_data = []
                        for row in new_data:
                            # 若数据长度大于表头长度，截取与表头长度相同的数据
                            if len(row) > header_count:
                                final_data.append(row[:header_count])
                            # 若数据长度小于表头长度，用空字符串补齐
                            elif len(row) < header_count:
                                new_row = row + [''] * (header_count - len(row))
                                final_data.append(new_row)
                            else:
                                final_data.append(row)
                        df = pd.DataFrame(final_data, columns=headers)
                        df.to_excel(writer, sheet_name=sheet_name, index=False)

                # 获取工作表
                ws = writer.sheets[sheet_name]

                # 定义边框样式
                thin_border = Border(left=Side(style='thin'),
                                     right=Side(style='thin'),
                                     top=Side(style='thin'),
                                     bottom=Side(style='thin'))

                # 为表头单元格设置边框
                if multi_level:
                    for row in ws.iter_rows(max_row=num_levels):
                        for cell in row:
                            cell.border = thin_border
                else:
                    for cell in ws.iter_rows(min_row=1, max_row=1):
                        for c in cell:
                            c.border = thin_border

                # 为内容单元格设置边框
                start_row = num_levels + 1 if multi_level else 2
                if data is not None and len(data) > 0:
                    for row in ws.iter_rows(min_row=start_row):
                        for cell in row:
                            cell.border = thin_border
            output.seek(0)
            return output
        except Exception as e:
            print(f"错误：导出数据时发生错误 {e}")
            return None

    @staticmethod
    async def single_excel_import(excel_file_stream, multi_level=False):
        """
        从 Excel 文件导入数据
        :param file_path: Excel 文件的路径
        :param multi_level: 是否为多级表头
        :return: 数据和表头
        """
        try:
            # 读取文件内容
            file_content =await excel_file_stream.read()
            # 使用 BytesIO 将内容转换为文件对象
            file_obj = BytesIO(file_content)
            file_obj.seek(0)
            if multi_level:
                # 手动读取多级表头
                wb = load_workbook(file_obj)
                # 默认获取第一个工作表
                ws = wb.active
                single_headers = []
                for cell in ws[1]:
                    single_headers.append(cell.value)

                data = []
                for row in ws.iter_rows(min_row=2, values_only=True):
                    data_row = []
                    for cell in row:
                        data_row.append(cell if cell is not None else 0)
                    data.append(data_row)

                df = pd.DataFrame(data, columns=single_headers)
            else:
                # 默认读取第一个工作表
                df = pd.read_excel(file_obj)
                df = df.fillna(0)

            headers = df.columns.tolist()
            data = df.values.tolist()

            # 封装数据为数组对象
            result_data = []
            for row in data:
                row_dict = {}
                for i, header in enumerate(headers):
                    row_dict[header] = row[i]
                result_data.append(row_dict)
            await excel_file_stream.close()
            return result_data
        except FileNotFoundError:
            print(f"错误：未找到文件 {file_obj}")
            return []
        except Exception as e:
            print(f"错误：读取文件时发生未知错误 {e}")
            return []

if __name__ == "__main__":
    # 确保 D://demo 目录存在
    save_dir = 'D://demo'
    try:
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
    except OSError as e:
        print(f"错误：创建目录 {save_dir} 时发生错误 {e}")
        exit(1)

    # 多级表头示例
    multi_level_headers = [('序号', '序号'), ('个人信息', '名称'), ('个人信息', '年龄'), ('其他信息', '性别'), ('备注', '备注')]
    multi_level_data = [{"序号": 1,  "年龄": 25,'名称': '张三', '性别': '男','备注':'beizhu'},
                        {"序号": 2, '名称': '李四', "年龄": 30, '性别': '女','备注':'beizhu'},
                        {"序号": 3, '名称': '王五', "年龄": 30, '性别': '女'}]

    # 导出多级表头数据
    output = ExcelUtils.excel_export(multi_level_data, multi_level_headers, multi_level=True, sheet_name='多级表头')

    if output:
        # 保存文件
        file_path = os.path.join(save_dir, 'combined.xlsx')
        try:
            with open(file_path, 'wb') as f:
                f.write(output.getvalue())
            print(f"文件已成功保存到 {file_path}")
        except OSError as e:
            print(f"错误：保存文件时发生错误 {e}")
    else:
        print("未生成有效文件，不进行保存操作。")


    # 单个导入数据示例
    # try:
    #     multi_imported_data = ExcelUtils.single_excel_import(file_path, multi_level=True)
    #
    #     print("表头导入数据:", multi_imported_data)
    # except Exception as e:
    #     print(f"错误：导入数据时发生错误 {e}")
